pub mod error;
pub mod utils;
pub mod prelude;
mod websocket;
mod http;

use base64::engine::general_purpose;
use base64::Engine;
use sha1::{Digest, Sha1};

pub use error::*;
pub use http::*;

use tokio::net::{TcpListener, TcpStream};
use tokio::sync::{mpsc, RwLock};
use uuid::Uuid;
use websocket::{DefaultHandler, WebSocketConnectionPool, WebSocketHandler, WebSocketSession};
use std::sync::Arc;
use tracing::{debug, error, info, warn};
use crate::http::{HttpRequest, HttpResponse};
use crate::server::utils::__StaticHandler;
use crate::websockt::{WebSocketFrame, WebSocketPayload, WebSocketReadStream, WebSocketRole, WebSocketWriteStream};

pub struct Server {
    pub routers: Vec<Box<dyn HttpRouter + Send + Sync>>,
    pub states: HttpStateContainer,

    pub handler: Arc<RwLock<dyn WebSocketHandler + Send + Sync>>,
    pub pool: Arc<RwLock<WebSocketConnectionPool>>,
}

impl Server {
    pub fn new() -> Self {
        info!("Server is initializing with static file handler.");
        Self { 
            routers: vec![Box::new(__StaticHandler::new())],
            states: HttpStateContainer::new(),
            pool: Arc::new(RwLock::new(WebSocketConnectionPool::new())),
            handler: Arc::new(RwLock::new(DefaultHandler::new())),
        }
    }

    pub fn with_handler<H: WebSocketHandler + Send + Sync + 'static>(mut self, handler: H) -> Self {
        self.handler = Arc::new(RwLock::new(handler));
        self
    }

    pub fn mount<R: HttpRouter + Send + Sync + 'static>(mut self, router: R) -> Self {
        debug!("Mounting route: {:?}", router.route_info());
        self.routers.push(Box::new(router));
        self
    }

    pub fn manage<T: Send + Sync + 'static>(mut self, value: T) -> Self {
        debug!("Injecting managed state of type {}", std::any::type_name::<T>());
        self.states.insert(value);
        self
    }
}

impl Server {
    pub async fn run(self, addr: &str) -> Result<(), ServerError> {
        let listener = TcpListener::bind(addr).await?;
        info!("🚀 Server listening on {}", addr);
        
        let shared_routers = Arc::new(self.routers);
        let shared_states = Arc::new(self.states);
        
        loop {
            let (socket, peer_addr) = listener.accept().await?;
            info!("🌐 Accepted connection from {}", peer_addr);

            let routers = shared_routers.clone();
            let states = shared_states.clone();
            let pool = self.pool.clone();
            let handler = self.handler.clone();

            tokio::spawn(async move {
                if let Err(err) = Self::handle_connection(socket, routers, states, pool, handler).await {
                    error!("Connection handling error: {}", err);
                }
            });
        }
    }
}

impl Server {
    async fn handle_connection(
        mut stream: TcpStream, 

        routers: Arc<Vec<Box<dyn HttpRouter + Send + Sync>>>, 
        states: Arc<HttpStateContainer>,

        pool: Arc<RwLock<WebSocketConnectionPool>>,
        handler: Arc<RwLock<dyn WebSocketHandler + Send + Sync>>,
    ) -> Result<(), ServerError> {

        let mut websocket_mode = false;

        loop {
            let request = HttpRequest::receive_from(&mut stream).await?;
            
            let request = match request {
                Some(req) => req,
                None => {
                    info!("Connection closed by client");
                    return Ok(());
                }
            };

            if Self::is_websocket_request(&request) {
                match Self::accept_websocket(&mut stream, &request).await {
                    Ok(()) => {
                        info!("Entry websocket mode!!!");
                        websocket_mode = true;
                        break;
                    }
                    Err(err) => error!("Something error happend: {:?}", err.to_string()),
                }
                continue;
            }
    
            let connection_keep_alive = request.headers.is_keep_alive();
    
            // Handle request
            let mut handled = false;
            for router in routers.iter() {
                if let Some(mut response) = router.route(&request, request.line.method, &request.line.uri, &states).await {
                    info!("✅ Route matched for {} {}", request.line.method, request.line.uri);
                    handled = true;
                    response.headers.set_keep_alive(connection_keep_alive);
                    response.send_to(&mut stream).await?;
                    debug!("📤 Response sent with status: {}", response.line.status_code);
                    break;
                }
            }
    
            // No route to handle this request
            if !handled {
                warn!("⚠️ No route matched for {} {}", request.line.method, request.line.uri);
                let mut not_found = HttpResponse::not_found();
                not_found.headers.set_keep_alive(connection_keep_alive);
                not_found.send_to(&mut stream).await?;
            }
    
            if !connection_keep_alive {
                info!("🔌 Connection closed (no keep-alive)");
                break;
            }
        }

        if websocket_mode {
            Self::websocket_loop(stream, pool.clone(), handler).await?;
        }
    
        Ok(())
    }
}

impl Server {
    async fn websocket_loop(
        stream: TcpStream, 
        pool: Arc<RwLock<WebSocketConnectionPool>>,
        handler: Arc<RwLock<dyn WebSocketHandler + Send + Sync>>,
    ) -> Result<(), ServerError> {
        let id = Uuid::new_v4();

        // Add new clinet to pool
        let (tx, mut rx) = mpsc::unbounded_channel::<WebSocketFrame>();
        pool.write().await.register(id, tx.clone()).await;

        // Create read & write stream
        let (reader, writer) = stream.into_split();
        let mut read_stream = WebSocketReadStream::new(reader);
        let mut write_stream = WebSocketWriteStream::new(writer, WebSocketRole::Server);

        // Read frame from stream, and call handler 
        let handler_clone = handler.clone(); 
        let pool_clone = pool.clone();
        let read_task = tokio::spawn(async move {
            let server = WebSocketSession::new(id.clone(), pool_clone);
            let handler = handler_clone;
            handler.write().await.on_open(&id, &server).await;

            loop {
                let frame = read_stream.read_frame().await?;
                let frame = match frame {
                    Some(f) => f,
                    None => {
                        info!("Connection closed by client");
                        break;
                    },
                };

                match &frame.payload {
                    WebSocketPayload::Text(msg) => {
                        info!("🔌 WebSocket receive message");
                        handler.write().await.on_message(&id, &server, msg).await;
                    }
                    WebSocketPayload::Close => {
                        info!("🔌 WebSocket connection closing");
                        handler.write().await.on_close(&id, &server).await;
                        break;
                    }
                    WebSocketPayload::Ping => {
                        server.send_frame(WebSocketFrame::server_pong()).await;
                    }
                    WebSocketPayload::Pong => {}
                    _ => error!("Unsupported WebSocket opcode: {:?}", frame.payload),
                }
            }

            Ok::<(), ServerError>(())
        });

        // Listen the channel, send all frame back to client!
        let write_task = tokio::spawn(async move {
            while let Some(frame) = rx.recv().await {
                write_stream.send_frame(&frame).await?;
            }
            Ok::<(), ServerError>(())
        });

        let (result1, result2) 
            = tokio::try_join!(read_task, write_task).unwrap();
        result1?; result2?;

        pool.write().await.unregister(&id).await;

        Ok(())
    }

    pub async fn accept_websocket(
        stream: &mut TcpStream,
        request: &HttpRequest
    ) -> Result<(), ServerError> {
        let headers = &request.headers;

        // Check "Connection: Upgrade"
        if headers.connection.as_ref().map(|s| s.as_str()) != Some("Upgrade") {
            return Err(ServerError::UpgradeWebSocket("Missing or invalid 'Connection' header".into()));
        }

        // Check "Upgrade: websocket"
        if headers.upgrade.as_ref().map(|s| s.as_str()) != Some("websocket") {
            return Err(ServerError::UpgradeWebSocket("Missing or invalid 'Upgrade' header".into()));
        }

        // Check "Sec-WebSocket-Version: 13"
        if headers.sec_websocket_version.as_ref().map(|s| s.as_str()) != Some("13") {
            return Err(ServerError::UpgradeWebSocket("Unsupported WebSocket version".into()));
        }

        // Get "Sec-WebSocket-Key"
        let key = match headers.sec_websocket_key.as_ref() {
            Some(k) => k,
            None => return Err(ServerError::UpgradeWebSocket("Missing 'Sec-WebSocket-Key'".into())),
        };

        // Create Sec-WebSocket-Accept
        let mut hasher = Sha1::new();
        hasher.update(key.as_bytes());
        hasher.update(b"258EAFA5-E914-47DA-95CA-C5AB0DC85B11");
        let hash_result = hasher.finalize();
        let accept_key = general_purpose::STANDARD.encode(&hash_result);

        // Send Response
        let response = HttpResponse::builder()
            .status(101)
            .reason("Switching Protocols")
            .connection("Upgrade")
            .upgrade("websocket")
            .sec_websocket_accept(&accept_key)
            .build().unwrap();

        response.send_to(stream).await?;
        Ok(())
    }

    fn is_websocket_request(request: &HttpRequest) -> bool {
        // "/ws"
        request.line.uri.path.len() == 1 && request.line.uri.path[0] == "ws"
    }
}

